package simpletrace

import (
	"context"
	"encoding/json"
	"fmt"
	"sync/atomic"
	"time"
)

func Init(ctx context.Context, exporter Exporter, extraInfo map[string]any, logger Logger) {
	SetExtraInfoFn(extraInfo)
	SetLogger(logger)
	SetSpanExporter(exporter)
	exporter.Start(ctx)
}

func Close() {
	if spanExporter != nil {
		spanExporter.Close()
	}
}

const (
	DefaultMaxSpanCount = 300
)

type Span struct {
	TraceID  string `json:"trace_id"`
	SpanID   string `json:"span_id"`
	ParentID string `json:"parent_id,omitempty"`

	Name     string         `json:"name"`
	Pos      string         `json:"pos,omitempty"` // 函数调用位置，通常是文件名:行号
	Start    time.Time      `json:"start"`
	End      time.Time      `json:"end"`
	Args     map[string]any `json:"args,omitempty"`
	Returns  map[string]any `json:"returns,omitempty"`
	Children []*Span        `json:"children"`

	ExtraInfo map[string]any `json:"extra_info,omitempty"` // 额外信息

	HasError bool `json:"has_error,omitempty"` // 是否有错误

	IsAttached bool `json:"is_attached,omitempty"` // 是否连接回父节点展示

	parent *Span // internal only
}

type contextKey struct{}

var spanContextKey = contextKey{}

func WithSpanContext(ctx context.Context, sc *SpanContext) context.Context {
	return context.WithValue(ctx, spanContextKey, sc)
}

func SpanContextFrom(ctx context.Context) (*SpanContext, bool) {
	sc, ok := ctx.Value(spanContextKey).(*SpanContext)
	return sc, ok
}

type SpanContext struct {
	TraceID  string `json:"trace_id"`
	SpanID   string `json:"span_id"`
	ParentID string `json:"parent_id,omitempty"`
	ChildReq *atomic.Uint32
}

type SpanOption func(*Span)

func WithIDs(traceID, spanID, parentID string) SpanOption {
	return func(s *Span) {
		if traceID != "" {
			s.TraceID = traceID
		}
		if spanID != "" {
			s.SpanID = spanID
		}
		if parentID != "" {
			s.ParentID = parentID
		}
	}
}

func WithPos(pos string) SpanOption {
	return func(s *Span) {
		s.Pos = pos
	}
}

func WithExtraInfo(extraInfo map[string]any) SpanOption {
	return func(s *Span) {
		if extraInfo != nil {
			s.ExtraInfo = extraInfo
		}
	}
}

var extraInfoFn SpanOption

func SetExtraInfoFn(extraInfo map[string]any) {
	extraInfoFn = func(s *Span) {
		if extraInfo != nil {
			s.ExtraInfo = extraInfo
		}
	}
}

func StartSpan(ctx context.Context, name string, args map[string]any, isAttached bool, opts ...SpanOption) (*Span, context.Context) {

	spanCtx := &SpanContext{
		ChildReq: new(atomic.Uint32),
	}
	span := &Span{
		Name:       name,
		Start:      time.Now(),
		Args:       args,
		IsAttached: isAttached,
	}

	for _, opt := range append(opts, extraInfoFn) {
		if opt != nil {
			opt(span)
		}
	}

	if span.TraceID != "" || span.SpanID != "" {
		spanCtx.TraceID = span.TraceID
		if span.SpanID == "" {
			span.SpanID = "1"
		}
		spanCtx.SpanID = span.SpanID
		spanCtx.ParentID = span.ParentID
	} else {
		parentCtx, ok := SpanContextFrom(ctx)
		if ok {
			spanCtx.TraceID = parentCtx.TraceID
			spanCtx.SpanID = fmt.Sprintf("%s.%d", parentCtx.SpanID, parentCtx.ChildReq.Add(1))
			spanCtx.ParentID = parentCtx.SpanID
		} else {
			return nil, ctx
		}
	}

	span.TraceID = spanCtx.TraceID
	span.SpanID = spanCtx.SpanID
	span.ParentID = spanCtx.ParentID

	return span, WithSpanContext(ctx, spanCtx)
}

func (span *Span) Finish(ctx context.Context, returns map[string]interface{}, hasError ...bool) {
	if span == nil {
		return
	}

	span.End = time.Now()
	if returns != nil {
		span.Returns = returns
	}
	if len(hasError) > 0 {
		span.HasError = hasError[0]
	}

	if spanExporter != nil {
		spanExporter.Export(ctx, span)
	}
}

func (span *Span) String() (string, error) {
	jsonData, err := json.Marshal(span)
	str := string(jsonData)
	if err == nil && str != "" {
		return str, nil
	}
	span.Args = make(map[string]any)
	span.Returns = make(map[string]any)
	jsonData, err = json.Marshal(span)
	return string(jsonData), err
}

func (span *Span) Unmarshal(data []byte) error {
	if err := json.Unmarshal(data, span); err != nil {
		return fmt.Errorf("[strace] failed to unmarshal span: %w", err)
	}
	return nil
}

func (span *Span) IsNotSubTree() bool {
	return span == nil || span.SpanID == "" || span.SpanID == "1"
}

var spanExporter Exporter

func SetSpanExporter(exporter Exporter) {
	spanExporter = exporter
}

func buildTraceTree(spans []*Span) *Span {
	if len(spans) == 0 {
		return nil
	}

	spanMap := make(map[string]*Span)
	existSpan := make([]*Span, 0, len(spans))
	for _, span := range spans {
		if _, ok := spanMap[span.SpanID]; ok {
			continue
		}
		spanMap[span.SpanID] = span
		existSpan = append(existSpan, span)
	}

	var root *Span
	for _, span := range existSpan {
		if (span.ParentID == "" || span.ParentID == "0") && (span.SpanID == "1") {
			root = span // 找到根节点
			continue
		}
		parent, exists := spanMap[span.ParentID]
		if exists {
			parent.Children = append(parent.Children, span)
			span.parent = parent // 设置父节点引用
		}
	}

	return root
}

func findNodeBySpanID(node *Span, spanID string) *Span {
	if node.SpanID == spanID {
		return node
	}
	for _, child := range node.Children {
		if found := findNodeBySpanID(child, spanID); found != nil {
			return found
		}
	}
	return nil
}

// 递归找出耗时最长路径
func findLongestPath(node *Span) []*Span {
	path := []*Span{node}
	for len(node.Children) > 0 {
		var maxChild *Span
		var maxDuration time.Duration
		for _, child := range node.Children {
			dur := child.End.Sub(child.Start)
			if dur > maxDuration {
				maxDuration = dur
				maxChild = child
			}
		}
		node = maxChild
		path = append(path, node)
	}
	return path
}

func ParseTraceSpan(jsonStr string) (*Span, error) {
	var nodes = make([]FlatSpan, 0)
	err := json.Unmarshal([]byte(jsonStr), &nodes)
	if err != nil {
		return nil, err
	}
	return RestoreTraceContext(nodes), nil
}

type FlatSpan struct {
	TraceID    string                 `json:"trace_id"`
	SpanID     string                 `json:"span_id"`
	ParentID   string                 `json:"parent_id,omitempty"`
	Name       string                 `json:"name"`
	Pos        string                 `json:"pos,omitempty"` // 函数调用位置，通常是文件名:行号
	Start      time.Time              `json:"start"`
	End        time.Time              `json:"end"`
	Args       map[string]interface{} `json:"args,omitempty"`
	Returns    map[string]interface{} `json:"returns,omitempty"`
	HasError   bool                   `json:"has_error,omitempty"`   // 是否有错误
	IsAttached bool                   `json:"is_attached,omitempty"` // 是否连接回父节点展示
	ExtraInfo  map[string]interface{} `json:"extra_info,omitempty"`  // 额外信息
}

func FlattenTraceContext(span *Span) []FlatSpan {
	var flat []FlatSpan
	var dfs func(node *Span)
	dfs = func(node *Span) {
		if node == nil {
			return
		}
		flat = append(flat, FlatSpan{
			TraceID:    node.TraceID,
			SpanID:     node.SpanID,
			ParentID:   node.ParentID,
			Name:       node.Name,
			Pos:        node.Pos,
			Start:      node.Start,
			End:        node.End,
			Args:       node.Args,
			Returns:    node.Returns,
			HasError:   node.HasError,
			IsAttached: node.IsAttached,
			ExtraInfo:  node.ExtraInfo,
		})
		for _, child := range node.Children {
			dfs(child)
		}
	}
	dfs(span)
	return flat
}

func RestoreTraceContext(flat []FlatSpan) *Span {
	nodeMap := make(map[string]*Span)
	var root *Span

	// 创建所有节点
	for _, f := range flat {
		nodeMap[f.SpanID] = &Span{
			ExtraInfo:  f.ExtraInfo,
			TraceID:    f.TraceID,
			SpanID:     f.SpanID,
			ParentID:   f.ParentID,
			Name:       f.Name,
			Pos:        f.Pos,
			Start:      f.Start,
			End:        f.End,
			Args:       f.Args,
			Returns:    f.Returns,
			HasError:   f.HasError,
			IsAttached: f.IsAttached,
			Children:   []*Span{},
		}
	}

	// 关联父子
	for _, node := range nodeMap {
		if node.ParentID == "" {
			root = node
		} else if parent, ok := nodeMap[node.ParentID]; ok {
			parent.Children = append(parent.Children, node)
			node.parent = parent
		} else {
			root = node
		}
	}

	return root
}

type ctxTraceHeaderKey string

const (
	CtxHeaderSimpleTraceEnable   ctxTraceHeaderKey = "X-Simple-Trace-Enable"
	CtxHeaderSimpleTraceID       ctxTraceHeaderKey = "X-Simple-Trace-ID"
	CtxHeaderSimpleTraceParentID ctxTraceHeaderKey = "X-Simple-Trace-Parent-ID"
	CtxHeaderSimpleTraceSpanID   ctxTraceHeaderKey = "X-Simple-Trace-Span-ID"
)

func isEnableSimpleTraceByCtx(ctx context.Context) (context.Context, bool) {
	if ctx == nil {
		return ctx, false
	}
	enabled, ok := ctx.Value(CtxHeaderSimpleTraceEnable).(bool)
	if ok && enabled {
		return ctx, true
	}
	_, ok = SpanContextFrom(ctx)
	if ok {
		return ctx, true
	}
	return ctx, false
}

func IsEnable(ctx context.Context) bool {
	_, enable := isEnableSimpleTraceByCtx(ctx)
	return enable
}

func TraceCtxFrom(srcCtx context.Context, distCtx context.Context) context.Context {
	span, ok := SpanContextFrom(srcCtx)
	if !ok {
		return distCtx
	}

	_, ok = SpanContextFrom(distCtx)
	if ok {
		return distCtx
	}

	if span == nil || span.TraceID == "" || span.SpanID == "" {
		return distCtx
	}

	traceID, parentID, spanID := span.TraceID, span.SpanID, fmt.Sprintf("%s.%d", span.SpanID, 1)
	ctx := context.WithValue(distCtx, CtxHeaderSimpleTraceEnable, true)
	ctx = context.WithValue(ctx, CtxHeaderSimpleTraceID, traceID)
	ctx = context.WithValue(ctx, CtxHeaderSimpleTraceParentID, parentID)
	ctx = context.WithValue(ctx, CtxHeaderSimpleTraceSpanID, spanID)
	return ctx
}
